/*
  PubSubClient.cpp - A simple client for MQTT.
  Nick O'Leary
  http://knolleary.net
*/

#include "PubSubClient.h"
#include "Arduino.h"

// Suppress uninitialized member variable in all constructors because some memory can be saved with
// on-demand initialization of these members
// cppcheck-suppress uninitMemberVar
PubSubClient::PubSubClient()
{
	this->_state = MQTT_DISCONNECTED;
	this->_client = NULL;
	this->stream = NULL;
	setCallback(NULL);
}
// cppcheck-suppress uninitMemberVar
PubSubClient::PubSubClient(Client& client)
{
	this->_state = MQTT_DISCONNECTED;
	setClient(client);
	this->stream = NULL;
}
// cppcheck-suppress uninitMemberVar
PubSubClient::PubSubClient(IPAddress addr, uint16_t port, Client& client)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(addr, port);
	setClient(client);
	this->stream = NULL;
}
// cppcheck-suppress uninitMemberVar
PubSubClient::PubSubClient(IPAddress addr, uint16_t port, Client& client, Stream& stream)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(addr,port);
	setClient(client);
	setStream(stream);
}
// cppcheck-suppress uninitMemberVar
// cppcheck-suppress passedByValue
PubSubClient::PubSubClient(IPAddress addr, uint16_t port, MQTT_CALLBACK_SIGNATURE, Client& client)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(addr, port);
	setCallback(callback);
	setClient(client);
	this->stream = NULL;
}
// cppcheck-suppress uninitMemberVar
// cppcheck-suppress passedByValue
PubSubClient::PubSubClient(IPAddress addr, uint16_t port, MQTT_CALLBACK_SIGNATURE, Client& client,
                           Stream& stream)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(addr,port);
	setCallback(callback);
	setClient(client);
	setStream(stream);
}
// cppcheck-suppress uninitMemberVar
PubSubClient::PubSubClient(uint8_t *ip, uint16_t port, Client& client)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(ip, port);
	setClient(client);
	this->stream = NULL;
}
// cppcheck-suppress uninitMemberVar
PubSubClient::PubSubClient(uint8_t *ip, uint16_t port, Client& client, Stream& stream)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(ip,port);
	setClient(client);
	setStream(stream);
}
// cppcheck-suppress uninitMemberVar
// cppcheck-suppress passedByValue
PubSubClient::PubSubClient(uint8_t *ip, uint16_t port, MQTT_CALLBACK_SIGNATURE, Client& client)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(ip, port);
	setCallback(callback);
	setClient(client);
	this->stream = NULL;
}
// cppcheck-suppress uninitMemberVar
// cppcheck-suppress passedByValue
PubSubClient::PubSubClient(uint8_t *ip, uint16_t port, MQTT_CALLBACK_SIGNATURE, Client& client,
                           Stream& stream)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(ip,port);
	setCallback(callback);
	setClient(client);
	setStream(stream);
}
// cppcheck-suppress uninitMemberVar
PubSubClient::PubSubClient(const char* domain, uint16_t port, Client& client)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(domain,port);
	setClient(client);
	this->stream = NULL;
}
// cppcheck-suppress uninitMemberVar
PubSubClient::PubSubClient(const char* domain, uint16_t port, Client& client, Stream& stream)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(domain,port);
	setClient(client);
	setStream(stream);
}
// cppcheck-suppress uninitMemberVar
// cppcheck-suppress passedByValue
PubSubClient::PubSubClient(const char* domain, uint16_t port, MQTT_CALLBACK_SIGNATURE,
                           Client& client)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(domain,port);
	setCallback(callback);
	setClient(client);
	this->stream = NULL;
}
// cppcheck-suppress uninitMemberVar
// cppcheck-suppress passedByValue
PubSubClient::PubSubClient(const char* domain, uint16_t port, MQTT_CALLBACK_SIGNATURE,
                           Client& client, Stream& stream)
{
	this->_state = MQTT_DISCONNECTED;
	setServer(domain,port);
	setCallback(callback);
	setClient(client);
	setStream(stream);
}

bool PubSubClient::connect(const char *id)
{
	return connect(id,NULL,NULL,0,0,0,0,1);
}

bool PubSubClient::connect(const char *id, const char *user, const char *pass)
{
	return connect(id,user,pass,0,0,0,0,1);
}

bool PubSubClient::connect(const char *id, const char* willTopic, uint8_t willQos,
                           bool willRetain, const char* willMessage)
{
	return connect(id,NULL,NULL,willTopic,willQos,willRetain,willMessage,1);
}

bool PubSubClient::connect(const char *id, const char *user, const char *pass,
                           const char* willTopic, uint8_t willQos, bool willRetain, const char* willMessage)
{
	return connect(id,user,pass,willTopic,willQos,willRetain,willMessage,1);
}

bool PubSubClient::connect(const char *id, const char *user, const char *pass,
                           const char* willTopic, uint8_t willQos, bool willRetain, const char* willMessage,
                           bool cleanSession)
{
	if (!connected()) {
		int result = 0;

		if (domain != NULL) {
			result = _client->connect(this->domain, this->port);
		} else {
			result = _client->connect(this->ip, this->port);
		}
		if (result == 1) {
			nextMsgId = 1;
			// Leave room in the buffer for header and variable length field
			uint16_t length = MQTT_MAX_HEADER_SIZE;
			unsigned int j;

#if MQTT_VERSION == MQTT_VERSION_3_1
			uint8_t d[9] = {0x00,0x06,'M','Q','I','s','d','p', MQTT_VERSION};
#define MQTT_HEADER_VERSION_LENGTH 9
#elif MQTT_VERSION == MQTT_VERSION_3_1_1
			uint8_t d[7] = {0x00,0x04,'M','Q','T','T',MQTT_VERSION};
#define MQTT_HEADER_VERSION_LENGTH 7
#endif
			for (j = 0; j<MQTT_HEADER_VERSION_LENGTH; j++) {
				buffer[length++] = d[j];
			}

			uint8_t v;
			if (willTopic) {
				v = 0x04|(willQos<<3)|(willRetain<<5);
			} else {
				v = 0x00;
			}
			if (cleanSession) {
				v = v|0x02;
			}

			if(user != NULL) {
				v = v|0x80;

				if(pass != NULL) {
					v = v|(0x80>>1);
				}
			}

			buffer[length++] = v;

			buffer[length++] = ((MQTT_KEEPALIVE) >> 8);
			buffer[length++] = ((MQTT_KEEPALIVE) & 0xFF);

			CHECK_STRING_LENGTH(length,id)
			length = writeString(id,buffer,length);
			if (willTopic) {
				CHECK_STRING_LENGTH(length,willTopic)
				length = writeString(willTopic,buffer,length);
				CHECK_STRING_LENGTH(length,willMessage)
				length = writeString(willMessage,buffer,length);
			}

			if(user != NULL) {
				CHECK_STRING_LENGTH(length,user)
				length = writeString(user,buffer,length);
				if(pass != NULL) {
					CHECK_STRING_LENGTH(length,pass)
					length = writeString(pass,buffer,length);
				}
			}

			write(MQTTCONNECT,buffer,length-MQTT_MAX_HEADER_SIZE);

			lastInActivity = lastOutActivity = millis();

			while (!_client->available()) {
				unsigned long t = millis();
				if (t-lastInActivity >= ((int32_t) MQTT_SOCKET_TIMEOUT*1000UL)) {
					_state = MQTT_CONNECTION_TIMEOUT;
					_client->stop();
					return false;
				}
			}
			uint8_t llen;
			uint16_t len = readPacket(&llen);

			if (len == 4) {
				if (buffer[3] == 0) {
					lastInActivity = millis();
					pingOutstanding = false;
					_state = MQTT_CONNECTED;
					return true;
				} else {
					_state = buffer[3];
				}
			}
			_client->stop();
		} else {
			_state = MQTT_CONNECT_FAILED;
		}
		return false;
	}
	return true;
}

// reads a byte into result
bool PubSubClient::readByte(uint8_t * result)
{
	uint32_t previousMillis = millis();
	while(!_client->available()) {
		yield();
		uint32_t currentMillis = millis();
		if(currentMillis - previousMillis >= ((int32_t) MQTT_SOCKET_TIMEOUT * 1000)) {
			return false;
		}
	}
	*result = _client->read();
	return true;
}

// reads a byte into result[*index] and increments index
bool PubSubClient::readByte(uint8_t * result, uint16_t * index)
{
	uint16_t current_index = *index;
	uint8_t * write_address = &(result[current_index]);
	if(readByte(write_address)) {
		*index = current_index + 1;
		return true;
	}
	return false;
}

uint16_t PubSubClient::readPacket(uint8_t* lengthLength)
{
	uint16_t len = 0;
	if(!readByte(buffer, &len)) {
		return 0;
	}
	bool isPublish = (buffer[0]&0xF0) == MQTTPUBLISH;
	uint32_t multiplier = 1;
	uint16_t length = 0;
	uint8_t digit = 0;
	uint16_t skip = 0;
	uint8_t start = 0;

	do {
		if (len == 5) {
			// Invalid remaining length encoding - kill the connection
			_state = MQTT_DISCONNECTED;
			_client->stop();
			return 0;
		}
		if(!readByte(&digit)) {
			return 0;
		}
		buffer[len++] = digit;
		length += (digit & 127) * multiplier;
		multiplier *= 128;
	} while ((digit & 128) != 0);
	*lengthLength = len-1;

	if (isPublish) {
		// Read in topic length to calculate bytes to skip over for Stream writing
		if(!readByte(buffer, &len)) {
			return 0;
		}
		if(!readByte(buffer, &len)) {
			return 0;
		}
		skip = (buffer[*lengthLength+1]<<8)+buffer[*lengthLength+2];
		start = 2;
		if (buffer[0]&MQTTQOS1) {
			// skip message id
			skip += 2;
		}
	}

	for (uint16_t i = start; i<length; i++) {
		if(!readByte(&digit)) {
			return 0;
		}
		if (this->stream) {
			if (isPublish && len-*lengthLength-2>skip) {
				this->stream->write(digit);
			}
		}
		if (len < MQTT_MAX_PACKET_SIZE) {
			buffer[len] = digit;
		}
		len++;
	}

	if (!this->stream && len > MQTT_MAX_PACKET_SIZE) {
		len = 0; // This will cause the packet to be ignored.
	}

	return len;
}

bool PubSubClient::loop()
{
	if (connected()) {
		unsigned long t = millis();
		if ((t - lastInActivity > MQTT_KEEPALIVE*1000UL) || (t - lastOutActivity > MQTT_KEEPALIVE*1000UL)) {
			if (pingOutstanding) {
				this->_state = MQTT_CONNECTION_TIMEOUT;
				_client->stop();
				return false;
			} else {
				buffer[0] = MQTTPINGREQ;
				buffer[1] = 0;
				_client->write(buffer,2);
				lastOutActivity = t;
				lastInActivity = t;
				pingOutstanding = true;
			}
		}
		if (_client->available()) {
			uint8_t llen;
			uint16_t len = readPacket(&llen);
			if (len > 0) {
				lastInActivity = t;
				uint8_t type = buffer[0]&0xF0;
				if (type == MQTTPUBLISH) {
					if (callback) {
						uint16_t tl = (buffer[llen+1]<<8)+buffer[llen+2]; /* topic length in bytes */
						memmove(buffer+llen+2,buffer+llen+3,tl); /* move topic inside buffer 1 byte to front */
						buffer[llen+2+tl] = 0; /* end the topic as a 'C' string with \x00 */
						char *topic = (char*) buffer+llen+2;
						uint8_t *payload;
						// msgId only present for QOS>0
						if ((buffer[0]&0x06) == MQTTQOS1) {
							uint16_t msgId = 0;
							msgId = (buffer[llen+3+tl]<<8)+buffer[llen+3+tl+1];
							payload = buffer+llen+3+tl+2;
							callback(topic,payload,len-llen-3-tl-2);

							buffer[0] = MQTTPUBACK;
							buffer[1] = 2;
							buffer[2] = (msgId >> 8);
							buffer[3] = (msgId & 0xFF);
							_client->write(buffer,4);
							lastOutActivity = t;

						} else {
							payload = buffer+llen+3+tl;
							callback(topic,payload,len-llen-3-tl);
						}
					}
				} else if (type == MQTTPINGREQ) {
					buffer[0] = MQTTPINGRESP;
					buffer[1] = 0;
					_client->write(buffer,2);
				} else if (type == MQTTPINGRESP) {
					pingOutstanding = false;
				}
			} else if (!connected()) {
				// readPacket has closed the connection
				return false;
			}
		}
		return true;
	}
	return false;
}

bool PubSubClient::publish(const char* topic, const char* payload)
{
	return publish(topic,(const uint8_t*)payload,strlen(payload),false);
}

bool PubSubClient::publish(const char* topic, const char* payload, bool retained)
{
	return publish(topic,(const uint8_t*)payload,strlen(payload),retained);
}

bool PubSubClient::publish(const char* topic, const uint8_t* payload, unsigned int plength)
{
	return publish(topic, payload, plength, false);
}

bool PubSubClient::publish(const char* topic, const uint8_t* payload, unsigned int plength,
                           bool retained)
{
	if (connected()) {
		if (MQTT_MAX_PACKET_SIZE < MQTT_MAX_HEADER_SIZE + 2+strlen(topic) + plength) {
			// Too long
			return false;
		}
		// Leave room in the buffer for header and variable length field
		uint16_t length = MQTT_MAX_HEADER_SIZE;
		length = writeString(topic,buffer,length);
		uint16_t i;
		for (i=0; i<plength; i++) {
			buffer[length++] = payload[i];
		}
		uint8_t header = MQTTPUBLISH;
		if (retained) {
			header |= 1;
		}
		return write(header,buffer,length-MQTT_MAX_HEADER_SIZE);
	}
	return false;
}

bool PubSubClient::publish_P(const char* topic, const char* payload, bool retained)
{
	return publish_P(topic, (const uint8_t*)payload, strlen(payload), retained);
}

bool PubSubClient::publish_P(const char* topic, const uint8_t* payload, unsigned int plength,
                             bool retained)
{
	unsigned int rc = 0;
	uint16_t tlen;
	unsigned int pos = 0;
	unsigned int i;
	uint8_t header;
	unsigned int len;

	if (!connected()) {
		return false;
	}

	tlen = strlen(topic);

	header = MQTTPUBLISH;
	if (retained) {
		header |= 1;
	}
	buffer[pos++] = header;
	len = plength + 2 + tlen;
	do {
		uint8_t digit;
		digit = len % 128;
		len = len / 128;
		if (len > 0) {
			digit |= 0x80;
		}
		buffer[pos++] = digit;
	} while(len>0);

	pos = writeString(topic,buffer,pos);

	rc += _client->write(buffer,pos);

	for (i=0; i<plength; i++) {
		rc += _client->write((char)pgm_read_byte_near(payload + i));
	}

	lastOutActivity = millis();

	return rc == tlen + 4 + plength;
}

bool PubSubClient::beginPublish(const char* topic, unsigned int plength, bool retained)
{
	if (connected()) {
		// Send the header and variable length field
		uint16_t length = MQTT_MAX_HEADER_SIZE;
		length = writeString(topic,buffer,length);
		uint8_t header = MQTTPUBLISH;
		if (retained) {
			header |= 1;
		}
		size_t hlen = buildHeader(header, buffer, plength+length-MQTT_MAX_HEADER_SIZE);
		uint16_t rc = _client->write(buffer+(MQTT_MAX_HEADER_SIZE-hlen),length-(MQTT_MAX_HEADER_SIZE-hlen));
		lastOutActivity = millis();
		return (rc == (length-(MQTT_MAX_HEADER_SIZE-hlen)));
	}
	return false;
}

int PubSubClient::endPublish()
{
	return 1;
}

size_t PubSubClient::write(uint8_t data)
{
	lastOutActivity = millis();
	return _client->write(data);
}

size_t PubSubClient::write(const uint8_t *buffer, size_t size)
{
	lastOutActivity = millis();
	return _client->write(buffer,size);
}

size_t PubSubClient::buildHeader(uint8_t header, uint8_t* buf, uint16_t length)
{
	uint8_t lenBuf[4];
	uint8_t llen = 0;
	uint8_t pos = 0;
	uint16_t len = length;
	do {
		uint8_t digit;
		digit = len % 128;
		len = len / 128;
		if (len > 0) {
			digit |= 0x80;
		}
		lenBuf[pos++] = digit;
		llen++;
	} while(len>0);

	buf[4-llen] = header;
	for (int i=0; i<llen; i++) {
		buf[MQTT_MAX_HEADER_SIZE-llen+i] = lenBuf[i];
	}
	return llen+1; // Full header size is variable length bit plus the 1-byte fixed header
}

bool PubSubClient::write(uint8_t header, uint8_t* buf, uint16_t length)
{
	uint16_t rc;
	uint8_t hlen = buildHeader(header, buf, length);

#ifdef MQTT_MAX_TRANSFER_SIZE
	uint8_t* writeBuf = buf+(MQTT_MAX_HEADER_SIZE-hlen);
	uint16_t bytesRemaining = length+hlen;  //Match the length type
	uint8_t bytesToWrite;
	bool result = true;
	while((bytesRemaining > 0) && result) {
		bytesToWrite = (bytesRemaining > MQTT_MAX_TRANSFER_SIZE)?MQTT_MAX_TRANSFER_SIZE:bytesRemaining;
		rc = _client->write(writeBuf,bytesToWrite);
		result = (rc == bytesToWrite);
		bytesRemaining -= rc;
		writeBuf += rc;
	}
	return result;
#else
	rc = _client->write(buf+(MQTT_MAX_HEADER_SIZE-hlen),length+hlen);
	lastOutActivity = millis();
	return (rc == hlen+length);
#endif
}

bool PubSubClient::subscribe(const char* topic)
{
	return subscribe(topic, 0);
}

bool PubSubClient::subscribe(const char* topic, uint8_t qos)
{
	if (qos > 1) {
		return false;
	}
	if (MQTT_MAX_PACKET_SIZE < 9 + strlen(topic)) {
		// Too long
		return false;
	}
	if (connected()) {
		// Leave room in the buffer for header and variable length field
		uint16_t length = MQTT_MAX_HEADER_SIZE;
		nextMsgId++;
		if (nextMsgId == 0) {
			nextMsgId = 1;
		}
		buffer[length++] = (nextMsgId >> 8);
		buffer[length++] = (nextMsgId & 0xFF);
		length = writeString((char*)topic, buffer,length);
		buffer[length++] = qos;
		return write(MQTTSUBSCRIBE|MQTTQOS1,buffer,length-MQTT_MAX_HEADER_SIZE);
	}
	return false;
}

bool PubSubClient::unsubscribe(const char* topic)
{
	if (MQTT_MAX_PACKET_SIZE < 9 + strlen(topic)) {
		// Too long
		return false;
	}
	if (connected()) {
		uint16_t length = MQTT_MAX_HEADER_SIZE;
		nextMsgId++;
		if (nextMsgId == 0) {
			nextMsgId = 1;
		}
		buffer[length++] = (nextMsgId >> 8);
		buffer[length++] = (nextMsgId & 0xFF);
		length = writeString(topic, buffer,length);
		return write(MQTTUNSUBSCRIBE|MQTTQOS1,buffer,length-MQTT_MAX_HEADER_SIZE);
	}
	return false;
}

void PubSubClient::disconnect()
{
	buffer[0] = MQTTDISCONNECT;
	buffer[1] = 0;
	_client->write(buffer,2);
	_state = MQTT_DISCONNECTED;
	_client->flush();
	_client->stop();
	lastInActivity = lastOutActivity = millis();
}

uint16_t PubSubClient::writeString(const char* string, uint8_t* buf, uint16_t pos)
{
	const char* idp = string;
	uint16_t i = 0;
	pos += 2;
	while (*idp) {
		buf[pos++] = *idp++;
		i++;
	}
	buf[pos-i-2] = (i >> 8);
	buf[pos-i-1] = (i & 0xFF);
	return pos;
}


bool PubSubClient::connected()
{
	bool rc;
	if (_client == NULL ) {
		rc = false;
	} else {
		rc = (int)_client->connected();
		if (!rc) {
			if (this->_state == MQTT_CONNECTED) {
				this->_state = MQTT_CONNECTION_LOST;
				_client->flush();
				_client->stop();
			}
		}
	}
	return rc;
}

PubSubClient& PubSubClient::setServer(uint8_t * ip, uint16_t port)
{
	IPAddress addr(ip[0],ip[1],ip[2],ip[3]);
	return setServer(addr,port);
}

PubSubClient& PubSubClient::setServer(IPAddress ip, uint16_t port)
{
	this->ip = ip;
	this->port = port;
	this->domain = NULL;
	return *this;
}

PubSubClient& PubSubClient::setServer(const char * domain, uint16_t port)
{
	this->domain = domain;
	this->port = port;
	return *this;
}
// cppcheck-suppress passedByValue
PubSubClient& PubSubClient::setCallback(MQTT_CALLBACK_SIGNATURE)
{
	this->callback = callback;
	return *this;
}

PubSubClient& PubSubClient::setClient(Client& client)
{
	this->_client = &client;
	return *this;
}

PubSubClient& PubSubClient::setStream(Stream& stream)
{
	this->stream = &stream;
	return *this;
}

int PubSubClient::state()
{
	return this->_state;
}
